import torch
import torch.nn as nn
import torch.optim as optim
from model.TemperatureHumidityMLP import *

# 设置分类的类别数量
num_classes_temp = 42 - 30 + 1  # 温度范围30-42，总共13个类别
num_classes_humidity = 60 - 40 + 1  # 湿度范围40-60，总共21个类别

# 示例：初始化模型并进行前向传播
input_size = 7
hidden_size1 = 64
hidden_size2 = 64

model = TemperatureHumidityMLP(input_size, hidden_size1, hidden_size2, num_classes_temp, num_classes_humidity)

# 示例输入数据
input_data = torch.tensor([[22.5, 0.5, 1, 30, 3, 10, 25]], dtype=torch.float32)

# 前向传播，获取预测结果
temp_logits, humidity_logits = model(input_data)

# 输出结果
temp_class = torch.argmax(temp_logits, dim=1).item() + 30  # 将类别索引转换为温度值
humidity_class = torch.argmax(humidity_logits, dim=1).item() + 40  # 将类别索引转换为湿度值

print(f"Predicted Temperature: {temp_class}°C")
print(f"Predicted Humidity: {humidity_class}%")

#torch.save(model.state_dict(), 'ckpt/model.pth')
